#!/usr/bin/env python

import argparse
import platform
import subprocess
import psutil
import GPUtil
import distro
import sys
import os
import re

def get_os_release():
    # name, version, _ = distro.linux_distribution(full_distribution_name=False)
    version = distro.version()
    name = distro.name()
    return version + '-' + name

def get_HW_configuration():
    # get os name and version
    os_name = platform.system()
    os_version = get_os_release()

    # get cpu info
    arch = platform.machine()
    cpu_count = psutil.cpu_count(logical=False)

    cpu_info = subprocess.check_output(["cat", "/proc/cpuinfo"])
    cpu_model = [s.strip() for s in cpu_info.decode().split("\n") if "model name" in s][0]
    cpu_model = cpu_model.split(":")[1].strip()

    # get gpu info
    gpus = GPUtil.getGPUs()  # get all gpus
    gpu_list = []

    for gpu in gpus:
        gpu_id = gpu.id
        gpu_name = gpu.name
        gpu_total_memory = f"{gpu.memoryTotal}MB"
        gpu_list.append((gpu_id, gpu_name, gpu_total_memory))
    try:
        gpu_list = gpu_list + get_metaxgpu_info()
    except:
        pass


    # get memory
    mem_total = round(psutil.virtual_memory().total / (1024.0 **3), 2)

    info = {
        "os_name": os_name,
        "os_version": os_version,
        "arch": arch,     
        "cpu_count": cpu_count,
        "cpu_model": cpu_model,
        "mem_total": mem_total,
        "gpu_list": gpu_list,
        "maca_version": os.getenv('MACA_VERSION'),
    }

    return info

def parse_args():
  parser = argparse.ArgumentParser()
  parser.add_argument('--output', dest='output', type=str, default=None,
                        help="specify the path of the HW config file.")
  return parser.parse_args()


def get_metaxgpu_info():
    try:
        result = subprocess.run(['mx-smi','--show-memory'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        if result.returncode != 0:
            print("Error exectuing mx-smi:", result.stderr)
        output_str = result.stdout
        
        # Extract the number of GPUs
        gpu_count = len(re.findall(r"GPU#\d+", output_str))
        gpu_names = re.findall(r"GPU#\d+  ([A-Za-z0-9]+)", output_str)
        vram_totals = re.findall(r"vis_vram total\s*:\s*(\d+)", output_str)
        gpu_list = []
        for i in range(gpu_count):
            gpu_id = i
            gpu_name = gpu_names[i]
            gpu_total_memory = f"{int(vram_totals[i]) / 1024 }MB"
            gpu_list.append((gpu_id, gpu_name, gpu_total_memory))
        return gpu_list
    except Exception as err:
        print(err)
    

